package gsw.tool.datastructure.tree;

import java.util.NoSuchElementException;

import gsw.tool.common.ToolShow;

/**
 * 红黑树是每个节点都带有颜色属性的自平衡二叉查找树，颜色或红色或黑色。
 * 性质1:节点是红色或黑色；
 * 性质2:根节点是黑色。
 * 性质3:所有的NULL节点都为叶子节点，且颜色为空。
 * 性质4:每个红色节点的两个子节点都是黑色。(从每个叶子到根的所有路径上不能有两个连续的红色节点)
 * 性质5:从任一节点到其每个叶子的所有路径都包含相同数目的黑色节点。
 */
public class TreeRedBlack<E extends Comparable> extends TreeBinary<E> {
    static final boolean BLACK = true;
    static final boolean RED = false;
    //根节点
    Node<E> root;

    public static void test() {
        TreeRedBlack tree = new TreeRedBlack();
        tree.createData();
        tree.ergodicRecursion(tree.root);
    }

    /**
     * 添加红黑树节点：
     * 1.添加跟二叉搜索树的插入一样；
     * 2.修正位置
     *
     * @param data
     */
    public void put(E data) {
        Node<E> now = new Node<E>(data);
        if (root == null) {
            root = now;
            return;
        }

        Node<E> parent = root;
        // 像二叉查找树一样添加
        while (parent != null) {
            if (data.compareTo(parent.data) < 0) {
                if (parent.left == null) {
                    parent.left = now;
                    now.parent = parent;
                    break;
                } else {
                    parent = parent.left;
                }
            } else if (data.compareTo(parent.data) > 0) {
                if (parent.right == null) {
                    parent.right = now;
                    now.parent = parent;
                    break;
                } else {
                    parent = parent.right;
                }
            } else {
                return;
            }
        }
        //修正位置
        fixAfterInsertion(now);
    }

    /**
     * 添加后修正位置。
     * 将当前节点记作A，将A的父节点记作B，将B的兄弟节点记作C，将B的父结点记作D。
     * 修正过程如下：
     * 1. 将A设为红色，这样对整颗树的平衡影响较小；
     * 2. 如果A为根节点，则将根节点改成黑色即可，修正结束；
     * 3. 如果B是黑色，那么树就是平衡的，修正结束；
     * 4. 如果B是红色，并且B是D的左孩子，则判断C的颜色：
     * - 如果C是红色，则将B和C涂黑，D涂红，把当前节点(A)指向D，然后继续比较A;
     * - 如果C是黑色，并且A是B的左孩子，则把B涂黑，D涂红，然后对D右旋，然后继续比较A；
     * - 如果C是黑色，并且A是B的右孩子，则把当前节点(A)指向B,然后对新的A左旋。然后把B指向新A的父亲，把D指向新B的父亲，然后把B涂黑，D涂红，对D右旋，然后继续比较A。
     * 5. 如果B是红色，并且B是D的右孩子，则判断C的颜色：
     * - 如果C是红色，则将B和C涂黑，D涂红，把当前节点(A)指向D，然后继续比较A;
     * - 如果C是黑色，并且A是B的右孩子，则把B涂黑，D涂红，然后对D左旋，然后继续比较A；
     * - 如果C是黑色，并且A是B的左孩子，则把当前节点(A)指向B,然后对新的A右旋。然后把B指向新A的父亲，把D指向新B的父亲，然后把B涂黑，D涂红，对D左旋，然后继续比较A。
     *
     * @param node
     */
    private void fixAfterInsertion(Node<E> node) {
        node.color = RED;

        Node<E> parent = null;
        Node<E> brother = null;
        while (node != null && node != root && node.parent.color == RED) {
            parent = node.parent;
            if (parent.parent != null && parent.parent.left == parent) {
                brother = parent.parent.right;
                if (getColor(brother) == RED) {
                    setColor(parent, BLACK);
                    setColor(brother, BLACK);
                    setColor(parent.parent, RED);
                    node = parent.parent;
                } else {
                    if (parent.right == node) {
                        node = parent;
                        leftRoate(node);
                    }
                    parent = node.parent;
                    setColor(parent, BLACK);
                    setColor(parent.parent, RED);
                    rightRoate(parent.parent);
                }
            } else {
                brother = parent.parent == null ? null : parent.parent.left;
                if (getColor(brother) == RED) {
                    setColor(parent, BLACK);
                    setColor(brother, BLACK);
                    setColor(parent.parent, RED);
                    node = parent.parent;
                } else {
                    if (parent.left == node) {
                        node = parent;
                        rightRoate(node);
                    }
                    parent = node.parent;
                    setColor(parent, BLACK);
                    setColor(parent.parent, RED);
                    leftRoate(parent.parent);
                }
            }
        }

        root.color = BLACK;
    }

    /**
     * 删除元素
     * 1.像二叉搜索树一样删除
     * 2.修正位置
     *
     * @param data
     * @return
     */
    public Node<E> remove(E data) {
        Node<E> now = get(data);
        if (now == null) {
            throw new NoSuchElementException();
        }

        Node<E> p = get(data);
        if (p == null) {
            throw new NoSuchElementException();
        }
        //p的右子树的最左子节点r
        Node<E> r = nodeLeft(p.right);
        //p的左子树的最右子节点l
        Node<E> l = nodeRight(p.left);
        if (r != null) {
            p.data = r.data;
            //如果p的右子结点有左节点
            if (r != p.right) {
                r.parent.left = r.right;
            } else {
                p.right = p.right.right;
            }
            //如果被删节点是黑色，就对替换的节点进行修正,即从平衡破坏的地方开始向上修正
            if (now.color == BLACK) {
                fixAfterDelete(r.right);
            }
            r.left = r.right = r.parent = null;

        } else if (l != null) {
            p.data = l.data;
            //如果p的左子结点有右节点
            if (l != p.left) {
                l.parent.right = l.left;
            } else {
                p.left = p.left.left;
            }
            //如果被删节点是黑色，就对替换的节点进行修正,即从平衡破坏的地方开始向上修正
            if (now.color == BLACK) {
                fixAfterDelete(l.left);
            }
            l.left = l.right = l.parent = null;

            //如果p是叶子节点
        } else {
            //如果p是叶子节点，但不是根节点，且颜色为黑色就需要对其修正
            if (p.parent != null && p.color == BLACK) {
                fixAfterDelete(p);
            }

            if (p.parent == null) {
                root = null;
            } else if (p.parent.left == p) {
                p.parent.left = null;
            } else {
                p.parent.right = null;
            }
            p.parent = null;
        }
        return now;
    }

    /**
     * 查找元素：跟二叉查找树一样
     *
     * @param data
     * @return
     */
    public Node<E> get(E data) {
        if (root == null) {
            return null;
        }
        Node<E> now = root;
        while (now != null) {
            if (data.compareTo(now.data) < 0) {
                now = now.left;
            } else if (data.compareTo(now.data) > 0) {
                now = now.right;
            } else {
                return now;
            }
        }
        return null;
    }

    /**
     * 获取当前节点的最左节点
     *
     * @param data
     * @return
     */
    public Node<E> nodeLeft(Node<E> data) {
        if (data == null) {
            return null;
        }
        Node<E> now = data;
        while (now.left != null) {
            now = now.left;
        }
        return now;
    }

    /**
     * 获取当前节点的最右子节点
     *
     * @param data
     * @return
     */
    public Node<E> nodeRight(Node<E> data) {
        if (data == null) {
            return null;
        }
        Node<E> now = data;
        while (now.right != null) {
            now = now.right;
        }
        return now;
    }


    /**
     * 删除后修正位置。
     * TODO：待完善
     * 将当前节点记作A，将A的父节点记作B，将A的兄弟节点记作C，将B的父结点记作D。
     * 修正过程如下：
     * 1. 如果A为根节点，则将A设为黑色，修正结束；
     * 2.如果A为红色，则将A设为黑色，修正结束；
     * 3. 如果A为B的左子节点，则
     * - 先判断C的颜色，如果C为红色，则将C设为黑色，将B设为红色，然后对B左旋，然后将B指向新A的父亲，将C指向新B的兄弟
     * - 然后判断C的左右子节点是否都为黑色，如果都为黑色，则将C设为红色，将A指向B，然后继续判断A；
     * - 如果C的右子节点是黑色，C的左子节点是红色或黑色，则先将C的左子节点设为黑色，将C设为红色，对B右旋，将B指向新A的父亲，将C指向新B的兄弟；然后将C设为B的颜色，将B设为黑色，将B的右子节点设为黑色，再将B左旋，将A指向root节点。
     * - 如果C的右子节点是红色，C的左子节点是红色或黑色，则将C设为B的颜色，将B设为黑色，将B的右子节点设为黑色，再将B左旋，将A指向root节点。
     * 4. 如果A为B的右子节点，则
     * - 先判断C的颜色，如果C为红色，则将C设为黑色，将B设为红色，然后对B右旋，然后将B指向新A的父亲，将C指向新B的兄弟
     * - 然后判断C的左右子节点是否都为黑色，如果都为黑色，则将C设为红色，将A指向B，然后继续判断A；
     * - 如果C的左子节点是黑色，C的右子节点是红色，则先将C的右子节点设为黑色，将C设为红色，对B左旋，将B指向新A的父亲，将C指向新B的兄弟；然后将C设为B的颜色，将B设为黑色，将B的左子节点设为黑色，再将B右旋，将A指向root节点。
     * - 如果C的左子节点是红色，C的右子节点是红色或黑色，则将C设为B的颜色，将B设为黑色，将B的左子节点设为黑色，再将B右旋，将A指向root节点。
     *
     * @param node
     */
    private void fixAfterDelete(Node<E> node) {
        if (node == null) {
            return;
        }

        Node<E> parent = null;
        Node<E> brother = null;
        while (node != root && getColor(node) == BLACK) {
            parent = node.parent;
            if (node == parent.left) {
                brother = parent.right;

                if (getColor(brother) == RED) {
                    setColor(brother, BLACK);
                    setColor(parent, RED);
                    leftRoate(parent);
                    parent = node.parent;
                    brother = parent.right;
                }

                if (getColor(brother.left) == BLACK && getColor(brother.right) == BLACK) {
                    setColor(brother, RED);
                    node = node.parent;
                } else {
                    if (getColor(brother.right) == BLACK) {
                        setColor(brother.left, BLACK);
                        setColor(brother, RED);
                        rightRoate(brother);
                        parent = node.parent;
                        brother = parent.right;
                    }
                    setColor(brother, getColor(parent));
                    setColor(parent, BLACK);
                    setColor(brother.right, BLACK);
                    leftRoate(parent);
                    node = root;
                }
            } else {
                brother = parent.left;

                if (getColor(brother) == RED) {
                    setColor(brother, BLACK);
                    setColor(parent, RED);
                    rightRoate(parent);
                    parent = node.parent;
                    brother = parent.left;
                }

                if (getColor(brother.left) == BLACK && getColor(brother.right) == BLACK) {
                    setColor(brother, RED);
                    node = node.parent;
                } else {
                    if (getColor(brother.left) == BLACK) {
                        setColor(brother.right, BLACK);
                        setColor(brother, RED);
                        leftRoate(brother);
                        parent = node.parent;
                        brother = parent.left;
                    }
                    setColor(brother, getColor(parent));
                    setColor(parent, BLACK);
                    setColor(brother.left, BLACK);
                    rightRoate(parent);
                    node = root;
                }
            }
        }

        setColor(node, BLACK);
    }

    /**
     * 设置颜色
     *
     * @param node
     */
    private void setColor(Node<E> node, boolean c) {
        if (node != null) {
            node.color = c;
        }
    }

    /**
     * 获取颜色
     *
     * @param node
     */
    private boolean getColor(Node<E> node) {
        return node == null ? BLACK : node.color;
    }

    /**
     * 左旋：把节点向左下移一位
     *
     * @param node
     */
    public void leftRoate(Node<E> node) {
        if (node == null) {
            return;
        }

        Node<E> right = node.right;
        // 1.把node的右子节点的左节点变成node的右子节点
        node.right = right.left;
        if (right.left != null) {
            right.left.parent = node;
        }

        // 2.把node的右子节点变出node的父节点的子节点
        right.parent = node.parent;
        if (node.parent == null) {
            root = right;
        } else if (node.parent.right == node) {
            node.parent.right = right;
        } else if (node.parent.left == node) {
            node.parent.left = right;
        }

        // 3.把node变成node的右子节点的左子节点
        right.left = node;
        node.parent = right;
    }

    /**
     * 右旋：把节点向右下移一位
     *
     * @param node
     */
    public void rightRoate(Node<E> node) {
        if (node == null) {
            return;
        }

        Node<E> left = node.left;
        // 1.把node的左子节点的右节点变成node的左子节点
        node.left = left.right;
        if (left.right != null) {
            left.right.parent = node;
        }

        // 2.把node的左子节点变出node的父节点的子节点
        left.parent = node.parent;
        if (node.parent == null) {
            root = left;
        } else if (node.parent.left == node) {
            node.parent.left = left;
        } else if (node.parent.right == node) {
            node.parent.right = left;
        }

        // 3.把node变成node的左子节点的右子节点
        left.right = node;
        node.parent = left;
    }

    /**
     * 递归前序遍历
     *
     * @param node 根节点
     */
    public void ergodicRecursion(Node<E> node) {
        if (node == null) {
            return;
        }
        ToolShow.log(node.toString());
        ergodicRecursion(node.left);
        ergodicRecursion(node.right);
    }

    /**
     * 测试数据
     */
    public void createData() {
        int[] list = new int[]{10, 1, 2, 44, 5, 7, 9};
        for (Integer i : list) {
            put((E) i);
        }
    }

    /**
     * 红黑树节点
     *
     * @param <E>
     */
    public static class Node<E> {
        E data;
        boolean color = BLACK;// 平衡因子
        Node<E> left;
        Node<E> right;
        Node<E> parent;

        public Node() {
        }

        public Node(E data) {
            super();
            this.data = data;
        }

        @Override
        public String toString() {
            StringBuffer sb = new StringBuffer();
            sb.append("Node{" + data + ", color=" + (color ? "黑" : "红"));
            if (left != null) {
                sb.append(", left=" + left.data);
            }
            if (right != null) {
                sb.append(", right=" + right.data);
            }
            if (parent != null) {
                sb.append(", parent=" + parent.data);
            }
            sb.append('}');
            return sb.toString();
        }
    }
}
